#!/usr/bin/env sh

set -e

: "${OVERRIDE:=false}"

: "${DEBUG:=false}"
[ "${DEBUG}" = "true" ] && set -x

: "${MODEL_TYPE:=yolov5s}"
: "${UPDATE_YOLOV5}"

BASEDIR="$(cd "$(dirname "$0")" && pwd)"

: "${GPU:=false}"

[ "${UPDATE_YOLOV5}" = "true" ] && [ -d "$BASEDIR/yolov5" ] && \rm -rf yolov5

main() {
    prev_dir="$(pwd)"
    cd "$BASEDIR"

    ! [ -d yolov5 ] &&
    mkdir -p yolov5 && cd yolov5 &&
    ! [ -d yolov5.git ] && git clone --bare https://github.com/ultralytics/yolov5.git &&
    git --git-dir=./yolov5.git --work-tree=. checkout HEAD -- data models utils LICENSE detect.py export.py val.py requirements.txt && \rm -rf yolov5.git && cd "${BASEDIR}"

    exists=0

    ! [ -d "${BASEDIR}/model_repository/torchscript_${MODEL_TYPE}/1" ] && mkdir -p "${BASEDIR}/model_repository/torchscript_${MODEL_TYPE}/1"
    [ -f "${BASEDIR}/${MODEL_TYPE}.torchscript" ] && exists=1
    if [ "${exists}" -eq 0 ] || [ "$OVERRIDE" = "true" ]; then
		if [ "$GPU" = "true" ]; then
			python3 yolov5/export.py --weights "$MODEL_TYPE.pt" --include torchscript --batch-size "$1" --dynamic --device 0
		else
			python3 yolov5/export.py --weights "$MODEL_TYPE.pt" --include torchscript --batch-size "$1" --dynamic
		fi
    fi
    cp "${BASEDIR}/${MODEL_TYPE}.torchscript" "${BASEDIR}/model_repository/torchscript_${MODEL_TYPE}/1/model.pt"

    echo "Successfully export YOLOv5 weights for TorchScript."
    cd "${prev_dir}"
}

if ! [ "$#" -eq 1 ]; then
    echo "Usage: $0 <batch_size>. Set to 1 by default if not passed."
    main 1
else
    main "$@"
fi
